import cv2
import os
import random


def ensure_dir(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)


def crop_and_save(
    image_path, label_path, output_dir, class_names, per_class_max, image_counter
):
    image = cv2.imread(image_path)
    if image is None:
        print(f"Image not found or could not be loaded: {image_path}")
        return {}
    height, width = image.shape[:2]

    with open(label_path, "r") as file:
        lines = file.readlines()

    saved_count = {cname: 0 for cname in class_names}

    for index, line in enumerate(lines):
        parts = line.strip().split()
        class_id = int(parts[0])
        cname = class_names[class_id]

        if image_counter[cname] >= per_class_max[class_id]:
            continue

        x_center = float(parts[1]) * width
        y_center = float(parts[2]) * height
        box_width = float(parts[3]) * width
        box_height = float(parts[4]) * height

        x1 = int(x_center - box_width / 2)
        y1 = int(y_center - box_height / 2)
        x2 = int(x_center + box_width / 2)
        y2 = int(y_center + box_height / 2)

        crop_size = 128
        dx = random.randint(-int(box_width), int(box_width))
        dy = random.randint(-int(box_height), int(box_height))
        crop_x1 = max(int(x_center - crop_size / 2 + dx), 0)
        crop_y1 = max(int(y_center - crop_size / 2 + dy), 0)
        crop_x2 = min(int(x_center + crop_size / 2 + dx), width)
        crop_y2 = min(int(y_center + crop_size / 2 + dy), height)

        # 确保裁剪区域不包含其他目标
        overlap = False
        for other_line in lines:
            if other_line == line:
                continue
            other_parts = other_line.strip().split()
            other_x_center = float(other_parts[1]) * width
            other_y_center = float(other_parts[2]) * height
            other_box_width = float(other_parts[3]) * width
            other_box_height = float(other_parts[4]) * height

            other_x1 = int(other_x_center - other_box_width / 2)
            other_y1 = int(other_y_center - other_box_height / 2)
            other_x2 = int(other_x_center + other_box_width / 2)
            other_y2 = int(other_y_center + other_box_height / 2)

            if not (
                crop_x2 < other_x1
                or crop_x1 > other_x2
                or crop_y2 < other_y1
                or crop_y1 > other_y2
            ):
                overlap = True
                break

        if overlap:
            continue

        cropped_image = image[crop_y1:crop_y2, crop_x1:crop_x2]

        if cropped_image.size == 0:
            print(
                f"Cropped image is empty. Skipping... [Image: {image_path}, Label: {line}]"
            )
            continue

        intersect_w = max(0, min(x2, crop_x2) - max(x1, crop_x1))
        intersect_h = max(0, min(y2, crop_y2) - max(y1, crop_y1))
        intersect_area = intersect_w * intersect_h
        crop_area = (crop_x2 - crop_x1) * (crop_y2 - crop_y1)
        area_ratio = intersect_area / crop_area

        if (area_ratio < 0.13 and area_ratio > 0.03) or (
            area_ratio > 0.8 and area_ratio < 0.85
        ):
            class_dir = os.path.join(output_dir, cname)
            ensure_dir(class_dir)
            image_counter[cname] += 1  # 更新计数器
            filename = f"{cname}_{image_counter[cname]}.png"
            filepath = os.path.join(class_dir, filename)
            cv2.imwrite(filepath, cropped_image)
            print(f"Saved: {filepath}")


def process_dataset(dataset_dir, output_dir, class_names, per_class_max):
    images_dir = os.path.join(dataset_dir, "images")
    labels_dir = os.path.join(dataset_dir, "labels")
    image_files = [f for f in os.listdir(images_dir) if f.endswith(".jpg")]

    image_counter = {
        cname: 0 for cname in class_names
    }  # Initialize image counter for each class

    for filename in image_files:
        base_name = filename[:-4]
        image_path = os.path.join(images_dir, filename)
        label_path = os.path.join(labels_dir, base_name + ".txt")

        if os.path.exists(label_path):
            crop_and_save(
                image_path,
                label_path,
                output_dir,
                class_names,
                per_class_max,
                image_counter,
            )


# 示例使用
class_names = ["red", "blue", "yellow"]
per_class_max = [300, 300, 300]  # 每类的最大图像数
dataset_dir = "/home/hw/dataset/cone_total_new1"
output_dir = "/home/hw/dataset/homework_plus"

process_dataset(dataset_dir, output_dir, class_names, per_class_max)
